port sensor,display,pyb,image, ml, math, uos, gc
import time
from machine import UART
roi1=	[( 20,   210, 20, 20),
          ( 90,   210, 20, 20),
          ( 150,   210, 20, 20),
          ( 210,  210, 20, 20),
          ( 300,  210, 20, 20)]
black_THRESHOLD=(0, 30, -22, 23, -128, 80)
red_THRESHOLD=  (0, 80, 13, 127, -4, 127)
uart = UART(3, 9600)
lcd=display.SPIDisplay()
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.skip_frames(time=2000)
sensor.set_auto_whitebal(False)
sensor.set_auto_gain(False)
sensor.set_hmirror(True)
sensor.set_vflip(True)
flagx=1
flag1=1
flag2=1
flag3=0
flag4=0
net = None
labels = None
min_confidence = 0.96
try:
    net = ml.Model("trained.tflite", load_to_fb=uos.stat('trained.tflite')[6] > (gc.mem_free() - (64*1024)))
except Exception as e:
    raise Exception('Failed to load "trained.tflite", did you copy the .tflite and labels.txt file onto the mass-storage device? (' + str(e) + ')')
try:
    labels = [line.rstrip('\n') for line in open("labels.txt")]
except Exception as e:
    raise Exception('Failed to load "labels.txt", did you copy the .tflite and labels.txt file onto the mass-storage device? (' + str(e) + ')')
colors = [
    (255,   0,   0),
    (  0, 255,   0),
    (255, 255,   0),
    (  0,   0, 255),
    (255,   0, 255),
    (  0, 255, 255),
    (255, 255, 255),
]
threshold_list = [(math.ceil(min_confidence * 255), 255)]
def fomo_post_process(model, inputs, outputs):
    ob, oh, ow, oc = model.output_shape[0]
    x_scale = inputs[0].roi[2] / ow
    y_scale = inputs[0].roi[3] / oh
    scale = min(x_scale, y_scale)
    x_offset = ((inputs[0].roi[2] - (ow * scale)) / 2) + inputs[0].roi[0]
    y_offset = ((inputs[0].roi[3] - (ow * scale)) / 2) + inputs[0].roi[1]
    l = [[] for i in range(oc)]
    for i in range(oc):
        img = image.Image(outputs[0][0, :, :, i] * 255)
        blobs = img.find_blobs(
            threshold_list, x_stride=1, y_stride=1, area_threshold=1, pixels_threshold=1
        )
        for b in blobs:
            rect = b.rect()
            x, y, w, h = rect
            score = (
                img.get_statistics(thresholds=threshold_list, roi=rect).l_mean() / 255.0
            )
            x = int((x * scale) + x_offset)
            y = int((y * scale) + y_offset)
            w = int(w * scale)
            h = int(h * scale)
            l[i].append((x, y, w, h, score))
    return l
clock = time.clock()
while (True):
    lcd.write(sensor.snapshot())
    clock.tick()
    img = sensor.snapshot()
    uart.write("0")
    print(clock.fps(), "fps", end="\n\n")
    #for roi in roi1:
        #img.draw_rectangle(roi, color=(0, 0, 0), thickness=2)
    #if(flagx):
        #for r in img.find_rects(threshold=25000):
            #img.draw_rectangle(r.rect(), color=(0, 0, 255))
            #for p in r.corners():
                #img.draw_circle(p[0], p[1], 5, color=(0, 255, 0))
            #print("1")
            #pyb.delay(1000)
            #uart.write("1")
            #flag1=1
            #flagx=0
    if(flag1):
        if(flag2):
            data=0
            blob2=None
            blob3=None
            blob4=None
            flag = [0,0,0,0,0]
            img = sensor.snapshot().lens_corr(strength = 1.7 , zoom = 1.0)
            blob2 = img.find_blobs([black_THRESHOLD], roi=roi1[1])
            blob3 = img.find_blobs([black_THRESHOLD], roi=roi1[2])
            blob4 = img.find_blobs([black_THRESHOLD], roi=roi1[3])
            if blob2:
                flag[1] = 1
            if blob3:
                flag[2] = 1
            if blob4:
                flag[3] = 1
                if blob2:
                    if blob3:
                        if blob4:
                            pyb.delay(480)
                            uart.write("2")
                            print("2")
                            flag4=1
                            flag2=0
                            flag3=1

    if(flag4):
        clock.tick()
        img = sensor.snapshot()
        detection_made = False
        for i, detection_list in enumerate(net.predict([img], callback=fomo_post_process)):
            if i == 0: continue  # background class

            if len(detection_list) == 0: continue
            for x, y, w, h, score in detection_list:
                center_x = math.floor(x + (w / 2))
                center_y = math.floor(y + (h / 2))
            if y<100:# no detections for this class?
                detection_made = True
                print("********** %s **********" % labels[i])
                print(f"x {center_x}\ty {center_y}\tscore {score}")
                img.draw_circle((center_x, center_y, 12), color=colors[i])
        if detection_made:
            print("4")
            uart.write("4")
            #flag4=0
        print(clock.fps(), "fps", end="\n\n")

    if(flag3):
        blob22=None
        blob33=None
        blob44=None
        blob22 = img.find_blobs([red_THRESHOLD], roi=roi1[1])
        blob33 = img.find_blobs([red_THRESHOLD], roi=roi1[2])
        blob44 = img.find_blobs([red_THRESHOLD], roi=roi1[3])
        if blob22:
            flag[1] = 1
        if blob33:
            flag[2] = 1
        if blob44:
            flag[3] = 1
            if blob22:
                if blob33:
                    if blob44:
                            print("5")
                            pyb.delay(500)
                            uart.write("5")
